package com.tuniu.agents.impl.train;

import com.tuniu.agents.agent.Agent;
import com.tuniu.agents.agent.AgentManagerGroup;
import com.tuniu.agents.common.agent.SelectionAgent;
import com.tuniu.agents.common.i18n.I18nConstants;
import com.tuniu.agents.common.message.BaseAgentNewResp;
import com.tuniu.agents.common.message.CardTypeEnum;
import com.tuniu.agents.memory.MemoryManager;
import com.tuniu.agents.message.MessageLogger;
import com.tuniu.agents.message.MessageMemoryManager;
import com.tuniu.agents.message.MessageRender;
import com.tuniu.agents.message.MessageSender;
import com.tuniu.agents.tools.train.constants.TrainResponseConstants;
import com.tuniu.agents.tools.train.tool.*;
import com.tuniu.agents.util.ChatResponseUtils;
import com.tuniu.agents.util.I18nUtil;
import com.tuniu.agents.util.JsonUtil;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.util.Collections;
import java.util.List;
import java.util.Map;

import static com.tuniu.agents.tools.train.tool.SearchTrainDetailTool.convertDetailDataToMark;


@Slf4j
public class TrainTicketAgent extends SelectionAgent {

    private static final Logger logger = LoggerFactory.getLogger(TrainTicketAgent.class);

    private MemoryManager memoryManager;

    private MessageSender messageSender;

    private MessageRender messageRender;

    @Autowired
    private MessageMemoryManager messageMemoryManager;


    public TrainTicketAgent(AgentManagerGroup agentManagerGroup,
                            String agentId, MessageSender messageSender,String promptTemplateId,MessageLogger messageLogger, Map<String, String> toolAnalysisMap, String... tools) {
        super(agentManagerGroup, agentId, promptTemplateId,messageLogger,toolAnalysisMap, tools);
        this.messageSender = messageSender;
        this.memoryManager = agentManagerGroup.getMemoryManager();
        this.messageRender = agentManagerGroup.getMessageRender();
    }

    @Override
    public Flux<String> stream(String conversationId, List<Message> messages, Map<String, Object> context) {

        ChatResponse response = this.callResponse(conversationId, messages, context);

        String toolName = ChatResponseUtils.getFirstToolName(response);
        if (!StringUtils.hasText(toolName)) {
            return Flux.just(ChatResponseUtils.getFirstTextMessage(response));
        } else {
            String toolResponse = ChatResponseUtils.getFirstTextMessage(response);

            Agent recommendAgent = null;
            switch (toolName) {
                case SearchAllTrainTool.TOOLNAME, SearchConditionTrainTool.TOOLNAME:
                    if (!StringUtils.hasText(toolResponse) || toolResponse.equals("null")) {
                        log.error("TrainTicketAgent train agent get toolResponse error:{}", toolResponse);
                        //Return error message and store it in context
                        messageMemoryManager.send(conversationId,I18nUtil.getMessage(TrainResponseConstants.COMMON_TRAIN_LIST_ERROR),this.agentId,List.of(I18nUtil.getMessage(TrainResponseConstants.COMMON_TRAIN_LIST_ERROR)));
                        return Flux.empty();
                    }

                    recommendAgent = agentManager.getAgent(toolAnalysisMap.get(toolName));
                    break;
                case SearchTrainDetailTool.TOOLNAME:
                    if (!StringUtils.hasText(toolResponse) || toolResponse.equals("null")) {
                        log.error("TrainTicketAgent train agent get toolResponse error:{}", toolResponse);
                        messageMemoryManager.send(conversationId, I18nUtil.getMessage(TrainResponseConstants.COMMON_TRAIN_DETAIL_ERROR),this.agentId,List.of(I18nUtil.getMessage(TrainResponseConstants.COMMON_TRAIN_DETAIL_ERROR)));
                        return Flux.empty();
                    }
                    SearchTrainDetailTool.Response detailResponse = JsonUtil.fromJson(toolResponse, SearchTrainDetailTool.Response.class);

                    //For detail query, return card directly without model analysis and recommendation
                    //Get card information
                    SearchTrainDetailTool.TrainDetailData detailCardInfo = SearchTrainDetailTool.buildDetailCardInfo(detailResponse,conversationId);
                    if (detailCardInfo == null) {
                        log.error("TrainTicketAgent train agent get detailCardInfo error");
                        messageMemoryManager.send(conversationId,  I18nUtil.getMessage(I18nConstants.COMMON_SYSTEM_ERROR_MESSAGE_KEY),this.agentId,List.of( I18nUtil.getMessage(I18nConstants.COMMON_SYSTEM_ERROR_MESSAGE_KEY)));
                        return Flux.empty();
                    }
                    List<Object> cardInfo = Collections.singletonList(detailCardInfo);
                    String cardType = CardTypeEnum.TRAIN_DETAIL.getCardType();

                    //get detailCardInfoMap
                    Map<String, Object> detailCardInfoMap = JsonUtil.toMap(JsonUtil.toJson(detailResponse.data()));

                    //build markdownResponse with detailCardInfoMap
                    String acceptLanguage = I18nUtil.getLocaleLanguageTag(conversationId);
                    String markdownResponse = messageRender.renderTemplate("train/train-detail.ftlh",Map.of(I18nConstants.LANGUAGE_TAG, acceptLanguage), detailCardInfoMap);

                    //Extract basic information from train details and save as markdown in context
                    memoryManager.setMessageRight(conversationId, Collections.singletonList(new AssistantMessage(convertDetailDataToMark(detailResponse.data(),conversationId), Map.of("agentId", this.agentId))));

                    String result =  BaseAgentNewResp.sendCardMsg(
                            cardType,
                            cardInfo,
                            markdownResponse,
                            null
                    );

                    return Flux.just(result);
            }

            //List agent analysis and recommendation data
            if (recommendAgent == null) {
                log.error("TrainTicketAgent train agent get recommendAgent fail toolName:{}", toolName);
                messageMemoryManager.send(conversationId,  I18nUtil.getMessage(I18nConstants.COMMON_SYSTEM_ERROR_MESSAGE_KEY),this.agentId,List.of( I18nUtil.getMessage(I18nConstants.COMMON_SYSTEM_ERROR_MESSAGE_KEY)));
                return Flux.empty();
            }
            //List message recommendation
            return recommendAgent.stream(conversationId, toolResponse);
        }

    }

}
